#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2021 wanghch <wanghch@wanghch-pc>
#
# Distributed under terms of the MIT license.

"""

"""
import pandas as pd
import os
import numpy as np
from tqdm import tqdm
from talib.abstract import *
from feature_utils import *
import sys
import argparse
import datetime
import random
import json

today = datetime.datetime.today()
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('-e', '--export', type=str, default=None, help='export')
parser.add_argument('-n', '--num', type=int, default=5, help='export n days')
parser.add_argument('-s', '--strategy', type=str, default="bottom_top", help='strategy fun')

ARGS = parser.parse_args()
# SEQ_DAYS = [2 * i for i in range(5)]
SEQ_DAYS = [0, 2, 5, 7]
with open('feature_conf.json' , 'rb') as f:
    fea_conf = json.load(f)


data_dir = "ind/"
export_dir = "export_ind/"
out_dir = "out_ind/"
datafiles = os.listdir(data_dir)
date_range = [5, 13]

def export():
    if not os.path.exists(export_dir):
        os.makedirs(export_dir)

    for dfile in tqdm(datafiles):
        code, name = dfile.replace(".csv", "").split("_")
        abs_dfile = os.path.join(data_dir, dfile)
        df = pd.read_csv(abs_dfile)
        N = df.shape[0]
        df = add_extra_indicators(df)
        out = os.path.join(export_dir, dfile)
        fout = open(out, "w")
        fdata = []
        # for fn in FNAMES:
            # convert_to_id(df, fn)


        for index,row in df.iterrows():
            if index < N - ARGS.num:
                continue
            features = extract_feature(df, index)
            fstr = ",".join(list(map(str, features.tolist())))
            fout.write(fstr + "\n")
        fout.close()


def bottom_top(df, index, row):
    c1 = row.close
    h1 = row.high
    l1 = row.low
    #c5 = df.loc[index + 5, 'close']

    #l3 = df.loc[index + 3, 'low']
    #h3 = df.loc[index + 3, 'high']

    #dt = row.date
    #g1 = row.high
    #g2 = df.loc[index + 1, 'high']
    #g0 = df.loc[index - 1, 'high']


    #d1 = row.low
    #d2 = df.loc[index + 1, 'low']
    #d0 = df.loc[index - 1, 'low']

    tflag = True
    lflag = True
    target = None

    hcloses = []
    lcloses = []
    for sqd in range(1, 5):
        if df.loc[index + sqd, 'low'] > l1 and df.loc[index - sqd, 'low'] > l1:
            pass
        else:
            lflag = False
            break

    for sqd in range(1, 5):
        if df.loc[index + sqd, 'high'] < h1:
            pass
        else:
            tflag = False
            break

    lastc = df.loc[index + 9, 'close']


    if lflag == True \
            and (c1 * 1.06) < lastc :
        target = 1.0
    elif tflag == True \
            and (c1 * 0.96) > lastc \
            and random.random() < 0.3: # negative downsampling
        target = 0.0


    return target


def trend(df, index, row):
    c1 = row.close
    c3 = df.loc[index + 3 , 'close']
    c5 = df.loc[index + 5,  'close']

    l3 = df.loc[index + 3, 'low']
    h3 = df.loc[index + 3, 'high']

    dt = row.date
    g1 = row.high
    g2 = df.loc[index + 1, 'high']
    g0 = df.loc[index - 1, 'high']


    d1 = row.low
    d2 = df.loc[index + 1, 'low']
    d0 = df.loc[index - 1, 'low']

    tflag = False
    lflag = False
    target = None
    # down
    if c1 > c3 and c3 > c5 and (c1 * 0.94) > c5:
        target = 0.0
    if c1 < c3 and c3 < c5 and (c1 * 1.06) < c5:
        target = 1.0

    #if target is None and random.random() < 0.1:
    #    target = 0.0

    return target





def main():
    for dfile in tqdm(datafiles):
        code = dfile.replace(".csv", "")
        abs_dfile = os.path.join(data_dir, dfile)
        df = pd.read_csv(abs_dfile)
        N = df.shape[0]
        df = add_extra_indicators(df)
        out = os.path.join(out_dir, dfile)
        # for fn in FNAMES:
            # convert_to_id(df, fn)
        fout = open(out, "w")
        fdata = []
        for index,row in df.iterrows():
            if index < 100 or index >= N - 10:
                continue

            strategy_func = eval(ARGS.strategy)
            target = strategy_func(df, index, row)

            if not target is None:
                features = extract_feature(df, index)
                fstr = ",".join(list(map(str, features.tolist())))
                fout.write(fstr + "," + str(target) + "\n")
        fout.close()

def calc_rsi(df):
    features = []
    for d in [6, 12, 24]:
        dfn = 'rsi_' + str(d)
        features.append(dfn)
        df[dfn] = RSI(df, timeperiod = d) / 100
    return df, features

def calc_bbands(df):
    features = ['bbp', "b_up_low"]
    bands = BBANDS(df, timeperiod = 20, nbdevup = 2.0, nbdevdn = 2.0, matype = 0)
    df['bbp'] = (df.close - bands.lowerband) / (bands.upperband - bands.lowerband)
    df['b_up_low'] = bands.upperband / bands.lowerband
    return df, features

def calc_macd(df):
    macd_df = MACD(df, fastperiod=12, slowperiod=26, signalperiod=9)
    # macd  macdsignal  macdhist
    df = pd.concat([df, macd_df], axis = 1)
    macd_features = ['macd','macdsignal','macdhist']
    return df, macd_features

def calc_kdj(df):
    matype = 0
    # slowk      slowd
    stoch = STOCH(df, fastk_period=9, slowk_matype=matype, slowk_period=3, slowd_period=3, slowd_matype=matype)
    df['stoch_k'] = stoch['slowk'] / 100
    df['stoch_d'] = stoch['slowd'] / 100
    df['stoch_j'] = (stoch['slowk'] * 3 - stoch['slowd'] * 2)/100
    kdj_features = ['stoch_k', 'stoch_d', 'stoch_j']
    return df, kdj_features


def calc_close_ma_diff(df):
    sma_close_diff_names = []
    for d in date_range:
        diff_name = 'SMA_DIFF_' + str(d)
        sma_name = 'SMA_' + str(d)
        sma_close_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d)
        df[diff_name] = (df['close'] - df[sma_name]) / df[sma_name]

    return df, sma_close_diff_names


def calc_vol_ma_diff(df):
    sma_vol_diff_names = []
    for d in date_range:
        diff_name = 'VOL_SMA_DIFF_' + str(d)
        sma_name = 'VOL_SMA_' + str(d)
        sma_vol_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d, price = 'volume')
        df[diff_name] = (df['volume'] - df[sma_name]) / df[sma_name]
    return df, sma_vol_diff_names

def calc_stats(df):
    features = []
    features.append("BETA")
    df["BETA"] = BETA(df.high, df.low, timeperiod = 5)

    features.append("CORREL")
    df["CORREL"] = CORREL(df.high, df.low, timeperiod=30)

    # features.append("LINEARREG")
    # df["LINEARREG"] = LINEARREG(df.close, timeperiod=14)

    # features.append("LINEARREG_ANGLE")
    # df["LINEARREG_ANGLE"] = LINEARREG_ANGLE(df.close, timeperiod=14)

    # features.append("LINEARREG_INTERCEPT")
    # df["LINEARREG_INTERCEPT"] = LINEARREG_INTERCEPT(df.close, timeperiod=14)

    # features.append("LINEARREG_SLOPE")
    # df["LINEARREG_SLOPE"] = LINEARREG_SLOPE(df.close, timeperiod=14)

    features.append("STDDEV")
    df["STDDEV"] = STDDEV(df.close, timeperiod=5)

    # print(df[features].tail(10))

    return df, features



def calc_cdl(df):
    features = []
    for cdln in CDL_NAMES:
        cdln_func = eval(cdln)
        df[cdln] = (cdln_func(df) > 0).astype(int)
        features.append(cdln)
    return df, features

def calc_cci(df):
    df["CCI"] = CCI(df.high, df.low, df.close, timeperiod=14) / 100
    df["WILLR"] = WILLR(df.high, df.low, df.close, timeperiod=14) /100
    # print(df[["CCI", "WILLR"]].tail(10))
    return df, ["CCI", "WILLR"]

def convert_to_id(df, name):
    buckets = fea_conf[name]['buckets']
    labels = list(range(len(buckets) - 1))
    df[name + "_id"] = pd.cut(df[name], bins = buckets, labels = labels)
    df.drop(name,axis=1, inplace=True)
    df.rename(columns = {name + '_id': name}, inplace = True)

def add_extra_indicators(df):
    features = []

    indicator_funcs = [
        calc_rsi,
        calc_kdj,
        calc_bbands,
        calc_macd,
        calc_close_ma_diff,
        calc_vol_ma_diff,
        calc_stats,
        calc_cci,
        # calc_cdl
    ]
    for indif in indicator_funcs:
        df, fns = indif(df)
        features += fns
    #print(features)
    df = df.fillna(method="bfill").fillna(method="ffill").fillna(0)
    df.replace([np.inf, -np.inf], 0, inplace = True)
    return df

def extract_feature(df, index):
    # ma5 diff
    return df.loc[index, INF_COLUMNS]



if ARGS.export:
    export()
else:
    main()
